﻿#pragma once

#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include <map>
#include <string>
#include <vector>

namespace websocket {

enum OpCode : int {
    WS_CONTINUATION    = 0x0, // 连续帧
    WS_TEXT            = 0x1, // 文本帧
    WS_BINARY          = 0x2, // 二进制帧
    WS_RESERVED3       = 0x3, // 保留
    WS_RESERVED4       = 0x4,
    WS_RESERVED5       = 0x5,
    WS_RESERVED6       = 0x6,
    WS_RESERVED7       = 0x7, // 保留
    WS_CLOSE           = 0x8, // 关闭
    WS_PING            = 0x9, // PING
    WS_PONG            = 0xA, // POING
    WS_CTRL_RESERVED_B = 0xB, // 保留0xb-0xf
    WS_CTRL_RESERVED_C = 0xC,
    WS_CTRL_RESERVED_D = 0xD,
    WS_CTRL_RESERVED_E = 0xE,
    WS_CTRL_RESERVED_F = 0xF,
};

enum class Status : int {
    NORMAL                  = 1000,
    GOING_AWAY              = 1001,
    PROTOCOL_ERROR          = 1002,
    UNSUPPORTED_DATA        = 1003,
    RESERVED                = 1004,
    NO_STATUS               = 1005,
    ABNORMALLY_CLOSE        = 1006,
    INVALID_PAYLOAD         = 1007,
    POLICY_VIOLATION        = 1008,
    MESSAGE_TOO_BIG         = 1009,
    MANDATORY_REQUIRED      = 1010,
    INTERNAL_ENDPOINT_ERROR = 1011,
    SERVICE_RESTART         = 1012,
    TRY_AGAIN_LATER         = 1013,
    BAD_GATEWAY             = 1014,
    TLS_HANDSHAKE           = 1015,
};

#define WSH_HOST                 "Host"
#define WSH_UPGRADE              "Upgrade"
#define WSH_CONNECTION           "Connection"
#define WSH_ORIGIN               "Origin"
#define WSH_SEC_ACCEPT           "Sec-WebSocket-Accept"
#define WSH_SEC_KEY              "Sec-WebSocket-Key"
#define WSH_SEC_PROTOCOL         "Sec-WebSocket-Protocol"
#define WSH_SEC_VERSION          "Sec-WebSocket-Version"
#define WEBSOCKET_VERSION        "13"
#define WEBSOCKET_SUCCESS_STATUS "101"

class Handshaker {
public:
    static int CheckPackageLength(const void *pdata, uint32_t data_size);

    Handshaker(const char *pdata, uint32_t data_size);
    void ParsingHeaders();

    bool IsValidClientRequest();
    bool IsValidServerResponse(const std::string &sec_websocket_key,
                               const std::string &sec_websocket_protocol = std::string());

    std::string FindField(const std::string &key);

    // HTTP/1.1 101 Switching Protocols
    // GET /chat HTTP/1.1
    std::string FirstLinePart1();
    std::string FirstLinePart2();
    std::string FirstLinePart3();

    std::string FirstLineText();

private:
    std::string header_;
    std::string first_line_str_;
    std::vector<std::string> first_line_;
    std::map<std::string, std::string> fields_;
};

class FrameContext {
public:
    static int CheckFrameLength(const void *data, uint32_t size);

    FrameContext(uint8_t *pdata, uint32_t data_size);
    ~FrameContext();

    bool ProcessData();

    bool Fin() const;
    int OpCode() const;
    bool Mask() const;

    uint32_t PayloadSize();
    uint8_t *Payload();
    uint32_t MaskKey();
    uint32_t FrameLength();

private:
    uint8_t *origin_ptr_;
    uint32_t origin_size_;

    bool fin_;
    int opcode_;

    bool mask_;
    uint32_t payload_len_;
    uint8_t *payload_;
    uint8_t mask_key_[4];
    uint32_t data_offset_;
};
// 构建帧前缀，从FIN到Mask-key部分，不包含payload
std::string BuildFramePrefix(bool fin, bool mask, int opcode, uint32_t mask_key, uint32_t size);

// 合并帧，如果包含mask标志，payload将会加密
std::string BuildFrameWithPayload(const std::string &prefix, uint8_t *payload,
                                  uint32_t payload_size);
// 合并帧
std::string CombinePrefixAndPayload(const std::string &prefix, uint8_t *payload,
                                    uint32_t payload_size);
std::string BuildSingleFrame(bool fin, bool mask, int opcode, uint32_t mask_key, uint8_t *payload,
                             uint32_t payload_size);

// 客户端发给服务端的数据必需用4字节的随机整数加密
void MaskingPayload(const std::string &frame_prefix, uint8_t *payload, uint32_t payload_size);
void MaskingPayload(uint32_t mask_key, uint8_t *payload, uint32_t payload_size);

// 判断一个完整数据包的长度（可能是握手包或 Websocket帧）
int32_t CheckPackageLength(const void *data, uint32_t len);

// 获取关闭帧中携带的调试信息
uint16_t GetCloseInfo(FrameContext *ctx, std::string *debug = NULL);

namespace client {
// 握手请求
std::string BuildHandshakeRequest(std::string *sec_websocket_key, const std::string &uri,
                                  const std::string &host, const std::string &origin,
                                  const std::string &protocol);
// 帧数据全部经过mask_key加密
std::string BuildCloseFrame(Status status, const std::string &debug, uint32_t mask_key = 0);
std::string BuildPingFrame(const std::string &data, uint32_t mask_key = 0);
std::string BuildPongFrame(const std::string &data, uint32_t mask_key = 0);
std::string BuildSingleDataFrame(const std::string &data, bool text_or_binary,
                                 uint32_t mask_key = 0);
std::string BuildDataFrameHeader(const std::string &data, bool text_or_binary,
                                 uint32_t mask_key = 0);
std::string BuildDataFrameBody(const std::string &data, uint32_t mask_key = 0);
std::string BuildDataFrameTailer(const std::string &data, uint32_t mask_key = 0);
// 构建连续帧
void BuildConsecutiveFrames(std::vector<std::string> *frames, const std::string &data,
                            uint32_t max_data_size_of_single_frame, bool text_or_binary,
                            uint32_t mask_key);
} // namespace client
namespace server {
// 握手响应
std::string BuildHandshakeSuccessReply(const std::string &sec_websocket_key,
                                       const std::string &sec_websocket_protocol);

// status: 非 101 即认为握手失败
std::string BuildHandshakeFailedReply(int status, const std::string &desc,
                                      const std::map<std::string, std::string> &fields);
// 帧数据未经过mask_key加密
std::string BuildCloseFrame(Status status, const std::string &debug);
std::string BuildPingFrame(const std::string &data);
std::string BuildPongFrame(const std::string &data);
std::string BuildSingleDataFrame(const std::string &data, bool text_or_binary);
std::string BuildDataFrameHeader(const std::string &data, bool text_or_binary);
std::string BuildDataFrameBody(const std::string &data);
std::string BuildDataFrameTailer(const std::string &data);
// 构建连续帧
void BuildConsecutiveFrames(std::vector<std::string> *frames, const std::string &data,
                            uint32_t max_data_size_of_single_frame, bool text_or_binary);
} // namespace server
} // namespace websocket
